#include <chrono>
#include <iomanip>

#include "config.hpp"
#include "header.hpp"
#include "logger.hpp"
#include "muxer_mp4.hpp"
#include "utils.hpp"

/**
 * @brief 构造函数
 */
muxer_mp4::muxer_mp4()
{
    this->rbuf = nullptr;
    cmd = COMMAND_IDLE;
}

/**
 * @brief 构造函数，同时设置取 h264 码流的环形缓冲区
 * @param ringbuffer h264 码流环形缓冲区
 */
muxer_mp4::muxer_mp4(ringbuffer* ringbuffer, std::string ip)
{
    if (nullptr == ringbuffer) {
        return;
    }

    this->ip = ip;
    this->rbuf = ringbuffer;
}

/**
 * @brief 析构函数
 */
muxer_mp4::~muxer_mp4()
{
}

/**
 * @brief 设置环形缓冲区
 * @param ringbuffer h264 码流环形缓冲区
 */
int muxer_mp4::set(ringbuffer* ringbuffer, std::string ip)
{
    if (nullptr == ringbuffer) {
        return -1;
    }

    this->ip = ip;
    this->rbuf = ringbuffer;

    return 0;
}

/**
 * @brief 切片处理线程，根据 CUT POINT 对视频进行切片操作
 */
void muxer_mp4::thread()
{
    int ret = 0;
    int ridx = 0;
    int stat = STATE_IDLE;
    int64_t first_nalu_pts = 0;
    int64_t first_nalu_dts = 0;
    AVPacket packet = {};

    /* 预申请一段内存空间，用于读取 NALU 数据 */
    int max_size = 2 * 1024 * 1024;
    uint8_t* nalu = (uint8_t*)av_malloc(max_size);

    while (running) {
        if (STATE_IDLE == stat) {
            if (COMMAND_IDLE == cmd) {
                msleep(40);
            } else if (COMMAND_START == cmd) {
                header_t head;
                int count = 0; /* 向前找 IDR 帧计数 */

                cmd = COMMAND_IDLE;

                rbuf->lock();
                ridx = rbuf->llast_widx_get();
                while (true) {
                    int r = ridx;
                    ret = rbuf->readr(&r, (uint8_t*)&head, sizeof(header_t));
                    if (ret) {
                        Logger::error("Camera {} muxer, Read h264 nalu data failed, mp4 find idr", ip);
                        ridx = -1;
                        break;
                    }

                    if (head.sync[0] != 'H' || head.sync[1] != 'E' || head.sync[2] != 'A' || head.sync[3] != 'D') {
                        Logger::error("Camera {} muxer, Sync code not match, mp4 find idr", ip);
                        ridx = -1;
                        break;
                    }

                    if (1 == head.type) {
                        count++;
                        Logger::debug("Camera {} muxer, IDR {} found, index is {}", ip, count, ridx);

                        /* 向前找 2 帧 IDR 帧 */
                        if (count >= 2) {
                            break;
                        } else {
                            ridx = head.prev;
                        }
                    } else {
                        ridx = head.prev;
                    }
                }
                rbuf->unlock();

                if (-1 != ridx) {
                    std::string name = CUT_VIDEO_SAVE_PATH + ip + "/" + utils::time_get() + ".mp4";
                    ret = open(name);
                    if (ret) {
                        Logger::error("Camera {} muxer, Failed to open mp4 file", ip);
                    }
                    stat = STATE_WORK;
                }

                /* 记录当前 MP4 第一帧视频的时间戳 */
                first_nalu_pts = head.pts;
                first_nalu_dts = head.dts;

            } else if (COMMAND_STOP == cmd) {
                Logger::error("Camera {} muxer, There stop command should not appear in the idle state", ip);
                msleep(40);
            } else {
                msleep(40);
            }
        } else if (STATE_WORK == stat) {
            if (COMMAND_IDLE == cmd) {
                if (-1 == ridx) {
                    Logger::error("Camera {} muxer, ridx invalid, bypass write mp4", ip);
                    msleep(40);
                    continue;
                }

                header_t head;
                rbuf->lock();

                ret = rbuf->readr(&ridx, (uint8_t*)&head, sizeof(header_t));
                if (ret) {
                    rbuf->unlock();
                    msleep(40);
                    continue;
                }

                if (head.sync[0] != 'H' || head.sync[1] != 'E' || head.sync[2] != 'A' || head.sync[3] != 'D') {
                    rbuf->unlock();
                    Logger::error("Camera {} muxer, Sync code not match", ip);
                    msleep(40);
                    continue;
                }

                memset(nalu, 0, max_size);
                packet.data = nalu;
                packet.size = head.size;
                packet.flags = head.type;
                ret = rbuf->readr(&ridx, packet.data, packet.size);
                if (ret) {
                    rbuf->unlock();
                    av_free(packet.data);
                    msleep(40);
                    continue;
                }

                rbuf->unlock();

                /* MP4 文件的时间戳不使用 rtsp 时间戳，而是从 0 开始 */
                packet.pts = head.pts - first_nalu_pts;
                packet.dts = head.dts - first_nalu_dts;
                ret = write(&packet);
                if (ret) {
                    Logger::error("Camera {} muxer, Failed to write mp4 file", ip);
                    msleep(40);
                    continue;
                }

            }

            else if (COMMAND_START == cmd) {
                Logger::error("Camera {} muxer, The start command should not appear in the running state", ip);
                msleep(40);
            }

            else if (COMMAND_STOP == cmd) {
                cmd = COMMAND_IDLE;
                close();
                ridx = -1;
                stat = STATE_IDLE;
            } else {
                msleep(40);
            }
        }
    }

    /* 被 stop() 接口退出本线程，则主动关闭 MP4 文件，避免出现 .tmp 文件 */
    if (STATE_WORK == stat) {
        cmd = COMMAND_IDLE;
        close();
        ridx = -1;
        stat = STATE_IDLE;
    }

    av_free(packet.data);
}

/**
 * @brief 启动切片线程
 */
void muxer_mp4::start()
{
    if (!handle.joinable()) {
        running = true;
        handle = std::thread(&muxer_mp4::thread, this);
    }
}

/**
 * @brief 关闭切片线程
 */
void muxer_mp4::stop()
{
    if (handle.joinable()) {
        running = false;
        handle.join();
    }
}

/**
 * @brief 设置切片关键点
 */
int muxer_mp4::cut(command cmd, int64_t pts)
{
    this->cmd = cmd;
    this->pts = pts;
    return 0;
}

/**
 * @brief 创建 MP4 文件
 */
int muxer_mp4::open(std::string filename)
{
    this->filename = filename;
    std::string tempname = filename + ".tmp";

    /* 明确指定输出格式为 MP4 */
    AVOutputFormat* ofmt = av_guess_format("mp4", nullptr, nullptr);
    if (!ofmt) {
        Logger::error("Camera {} muxer, Could not find output format for MP4", ip);
        return -1;
    }

    if (avformat_alloc_output_context2(&fmtctx, ofmt, nullptr, tempname.c_str()) < 0) {
        Logger::error("Camera {} muxer, Could not create output context", ip);
        return -1;
    }

    codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    if (!codec) {
        Logger::error("Camera {} muxer, H.264 codec not found", ip);
        return -1;
    }

    stream = avformat_new_stream(fmtctx, codec);
    if (!stream) {
        Logger::error("Camera {} muxer, Failed to create video stream", ip);
        return -1;
    }

    AVCodecParameters* codecpar = stream->codecpar;
    codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
    codecpar->codec_id = AV_CODEC_ID_H264;
    codecpar->width = 1920;
    codecpar->height = 1080;
    codecpar->format = AV_PIX_FMT_YUV420P;

    if (!(fmtctx->oformat->flags & AVFMT_NOFILE)) {
        if (avio_open(&(fmtctx->pb), tempname.c_str(), AVIO_FLAG_WRITE) < 0) {
            Logger::error("Camera {} muxer, Could not open output file", ip);
            return -1;
        }
        Logger::info("Camera {} muxer, file create {}", ip, tempname);
    }

    if (avformat_write_header(fmtctx, nullptr) < 0) {
        Logger::error("Camera {} muxer, Error writing header to output file", ip);
        return -1;
    }

    return 0;
}

/**
 * @brief 写 MP4 文件
 */
int muxer_mp4::write(AVPacket* packet)
{
    if (av_interleaved_write_frame(fmtctx, packet) < 0) {
        Logger::error("Camera {} muxer, Error writing frame", ip);
        return -1;
    }

    return 0;
}

/**
 * @brief 关闭 MP4 文件
 */
int muxer_mp4::close()
{
    av_write_trailer(fmtctx);
    avio_closep(&(fmtctx->pb));

    /* tmp 文件重命名为 mp4 文件 */
    const char* finalname = filename.c_str();
    if (rename((filename + ".tmp").c_str(), finalname) != 0) {
        Logger::error("Camera {} muxer, Error renaming file to {}", ip, finalname);
        avformat_free_context(fmtctx);
        return -1;
    }

    Logger::info("Camera {} muxer, file rename {}", ip, finalname);
    avformat_free_context(fmtctx);

    return 0;
}

/**
 * @brief 毫秒延时
 */
void muxer_mp4::msleep(int ms)
{
    std::this_thread::sleep_for(std::chrono::milliseconds(ms));
}